# coding=utf-8
'''
此代码用于数据预处理
'''
import os
import cv2
import random
import numpy as np
import tensorflow as tf
import core.utils as utils
from core.config import cfg


class Dataset(object):
    """implement Dataset here"""

    # 定义初始化函数
    def __init__(self, dataset_type):
        # 根据数据类型选择注释文件
        self.annot_path = cfg.TRAIN.ANNOT_PATH if dataset_type == 'train' else cfg.TEST.ANNOT_PATH
        # 输入图像的大小，为[320, 352, 384, 416, 448, 480, 512, 544, 576, 608]中的一种
        self.input_sizes = cfg.TRAIN.INPUT_SIZE if dataset_type == 'train' else cfg.TEST.INPUT_SIZE
        # Batch_size,每次喂入的数据量
        self.batch_size = cfg.TRAIN.BATCH_SIZE if dataset_type == 'train' else cfg.TEST.BATCH_SIZE
        # 训练集使用图像增强，测试集不使用图像增强
        self.data_aug = cfg.TRAIN.DATA_AUG if dataset_type == 'train' else cfg.TEST.DATA_AUG
        # 训练数据大小[320, 352, 384, 416, 448, 480, 512, 544, 576, 608]
        self.train_input_sizes = cfg.TRAIN.INPUT_SIZE
        # 3种下采样方式，为[8,16,32]
        self.strides = np.array(cfg.YOLO.STRIDES)
        # 训练数据的类别
        self.classes = utils.read_class_names(cfg.YOLO.CLASSES)
        # 训练数据的类别总数
        self.num_classes = len(self.classes)
        # 加载anchor box
        self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS))
        # 每种缩放比例下，每个grid cell预测的anchor box数
        self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE
        # 每种尺度下的最大bounding box数
        self.max_bbox_per_scale = 150
        # 读取标注信息
        self.annotations = self.load_annotations(dataset_type)
        # 样本总量
        self.num_samples = len(self.annotations)
        # 每个EPOCH,需要喂入多少组数据
        self.num_batchs = int(np.ceil(self.num_samples / self.batch_size))
        # 当batch_count达到num_batchs代表训练了一个EPOCHS
        self.batch_count = 0

    # 读取标注文件
    def load_annotations(self, dataset_type):
        with open(self.annot_path, 'r') as f:
            txt = f.readlines()
            annotations = [line.strip() for line in txt if len(line.strip().split()[1:]) != 0]
        # 打乱注释内容的顺序
        np.random.shuffle(annotations)
        return annotations

    """
    如果一个类想被用于for...in 循环.类似于list和tuple那就必须实现__init__()方法
    该方法返回一个迭代器对象，此时需要迭代的对象是实例本身，所以返回是return self
    """
    def __iter__(self):
        return self

    """
    当该类实现__iter__()方法后，该类就是可迭代的
    使用for...in 循环时会不断调用__next__()拿到循环的下一个值
    直到遇到StopIteration错误时退出循环
    
    Returns:
        tuple类型 
    
    """
    def __next__(self):

        with tf.device('/gpu:0'):
            # 从输入尺寸中随机挑选一个数
            self.train_input_size = random.choice(self.train_input_sizes)
            # // 表示整数除法,返回不大于结果的一个最大的整数
            # /  表示浮点数除法
            # 若self.train_input_size为544则self.train_output_sizes=544//[8,16,32]=[68 34 17]
            self.train_output_sizes = self.train_input_size // self.strides
            # 图像打包：一个BATCH_SIZE的输入图像（每次喂入的图像大小，图像宽，图像高，图像维度）
            batch_image = np.zeros((self.batch_size, self.train_input_size, self.train_input_size, 3))
            # ground_truth
            batch_label_sbbox = np.zeros((self.batch_size, self.train_output_sizes[0], self.train_output_sizes[0],
                                          self.anchor_per_scale, 5 + self.num_classes))
            batch_label_mbbox = np.zeros((self.batch_size, self.train_output_sizes[1], self.train_output_sizes[1],
                                          self.anchor_per_scale, 5 + self.num_classes))
            batch_label_lbbox = np.zeros((self.batch_size, self.train_output_sizes[2], self.train_output_sizes[2],
                                          self.anchor_per_scale, 5 + self.num_classes))
            #预测选框
            batch_sbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
            batch_mbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
            batch_lbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))

            num = 0
            # 当训练过的batch数<一个epcho需要的batch数时：
            # 一个epcho，训练完一次所有样本
            if self.batch_count < self.num_batchs:
                # 当num<BATCH_SIZE时：
                while num < self.batch_size:
                    # index为样本元素的索引（第index个样本）
                    index = self.batch_count * self.batch_size + num
                    # 索引不能超出样本总量，超出就重新开始
                    if index >= self.num_samples: index -= self.num_samples
                    # 第index个样本的标注文本信息
                    annotation = self.annotations[index]
                    # 从标注文本信息解析标注信息
                    image, bboxes = self.parse_annotation(annotation)

                    label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = self.preprocess_true_boxes(
                        bboxes)
                    # 一个batch中第num个图像的信息
                    batch_image[num, :, :, :] = image
                    batch_label_sbbox[num, :, :, :, :] = label_sbbox
                    batch_label_mbbox[num, :, :, :, :] = label_mbbox
                    batch_label_lbbox[num, :, :, :, :] = label_lbbox
                    batch_sbboxes[num, :, :] = sbboxes
                    batch_mbboxes[num, :, :] = mbboxes
                    batch_lbboxes[num, :, :] = lbboxes
                    num += 1
                self.batch_count += 1
                return batch_image, batch_label_sbbox, batch_label_mbbox, batch_label_lbbox, \
                       batch_sbboxes, batch_mbboxes, batch_lbboxes
            else:
                self.batch_count = 0
                np.random.shuffle(self.annotations)
                raise StopIteration

    # 随机水平翻转
    def random_horizontal_flip(self, image, bboxes):

        if random.random() < 0.5:
            _, w, _ = image.shape
            image = image[:, ::-1, :]
            bboxes[:, [0, 2]] = w - bboxes[:, [2, 0]]

        return image, bboxes

    # 随机裁剪
    def random_crop(self, image, bboxes):

        if random.random() < 0.5:
            h, w, _ = image.shape
            max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)

            max_l_trans = max_bbox[0]
            max_u_trans = max_bbox[1]
            max_r_trans = w - max_bbox[2]
            max_d_trans = h - max_bbox[3]

            crop_xmin = max(0, int(max_bbox[0] - random.uniform(0, max_l_trans)))
            crop_ymin = max(0, int(max_bbox[1] - random.uniform(0, max_u_trans)))
            crop_xmax = max(w, int(max_bbox[2] + random.uniform(0, max_r_trans)))
            crop_ymax = max(h, int(max_bbox[3] + random.uniform(0, max_d_trans)))

            image = image[crop_ymin: crop_ymax, crop_xmin: crop_xmax]

            bboxes[:, [0, 2]] = bboxes[:, [0, 2]] - crop_xmin
            bboxes[:, [1, 3]] = bboxes[:, [1, 3]] - crop_ymin

        return image, bboxes

    # 随机翻转
    def random_translate(self, image, bboxes):

        if random.random() < 0.5:
            h, w, _ = image.shape
            max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)

            max_l_trans = max_bbox[0]
            max_u_trans = max_bbox[1]
            max_r_trans = w - max_bbox[2]
            max_d_trans = h - max_bbox[3]

            tx = random.uniform(-(max_l_trans - 1), (max_r_trans - 1))
            ty = random.uniform(-(max_u_trans - 1), (max_d_trans - 1))

            M = np.array([[1, 0, tx], [0, 1, ty]])
            image = cv2.warpAffine(image, M, (w, h))

            bboxes[:, [0, 2]] = bboxes[:, [0, 2]] + tx
            bboxes[:, [1, 3]] = bboxes[:, [1, 3]] + ty

        return image, bboxes


    """
    解析标注信息，返回图像增强后的图像和标签 
    Arguments:
        annotation:标注文本信息[image_path x_min,y_min,x_max,y_max,class_id]
    Returns:
        image:经过镜像、缩放、旋转且缩放为网络输入尺寸的图像
        bbox:ground_truth.经过上述处理后image的ground_truth
            [xmin, ymin, xmax, ymax,class_id]
    """
    def parse_annotation(self, annotation):
        #str.split()以空格分隔字符
        line = annotation.split()
        # 照片的路径
        image_path = line[0]
        if not os.path.exists(image_path):
            raise KeyError("%s does not exist ... " % image_path)
        # 读取图像
        image = np.array(cv2.imread(image_path))
        # bbox信息：xmin, ymin, xmax, ymax,class_index
        bboxes = np.array([list(map(lambda x: int(float(x)), box.split(','))) for box in line[1:]])
        # 图像增强
        if self.data_aug:
            # np.copy为深拷贝，原始对象改变时，拷贝对象不会改变
            # 水平镜像
            image, bboxes = self.random_horizontal_flip(np.copy(image), np.copy(bboxes))
            # 裁剪
            image, bboxes = self.random_crop(np.copy(image), np.copy(bboxes))
            # 旋转
            image, bboxes = self.random_translate(np.copy(image), np.copy(bboxes))
        # 将输入图像按照input_size缩放，缩放后尺寸不足input_size的部分用中性灰填充
        image, bboxes = utils.image_preporcess(np.copy(image), [self.train_input_size, self.train_input_size],
                                               np.copy(bboxes))
        return image, bboxes

    # 计算IOU值
    def bbox_iou(self, boxes1, boxes2):
        # 此处使用...的秒处
        # boxes1可能包含一组[b_x, b_y, b_w, b_h]，也可能包含多组
        # 而a=np.array([])为一维数组不能使用a[:,2:]只能当数组维数>=2时才可以使用
        # 但一维数组可以使用a[...,2:]
        # 可以假设boxes = [b_x, b_y, b_w, b_h]
        boxes1 = np.array(boxes1)
        boxes2 = np.array(boxes2)
        # 计算boxes的面积：b_w*b_h
        boxes1_area = boxes1[..., 2] * boxes1[..., 3]
        boxes2_area = boxes2[..., 2] * boxes2[..., 3]
        # boxes1=boxes1的[左上角坐标，右下角坐标][x1,y1,x2,y2]
        boxes1 = np.concatenate([boxes1[..., :2] - boxes1[..., 2:] * 0.5,
                                 boxes1[..., :2] + boxes1[..., 2:] * 0.5], axis=-1)
        # boxes2=boxes2的[左上角坐标，右下角坐标]
        boxes2 = np.concatenate([boxes2[..., :2] - boxes2[..., 2:] * 0.5,
                                 boxes2[..., :2] + boxes2[..., 2:] * 0.5], axis=-1)
        # 计算交集部分左上角坐标
        left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
        # 计算交集部分右下角坐标
        right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])
        # 交集部分的w,h
        inter_section = np.maximum(right_down - left_up, 0.0)
        # 计算交集面积
        inter_area = inter_section[..., 0] * inter_section[..., 1]
        # 计算并集面积
        union_area = boxes1_area + boxes2_area - inter_area
        # 返回IOU值
        return inter_area / union_area


    """
    计算3个先验框和真实框的iou值，筛选iou值>0.3的先验框并标记索引， 
    然后将真实框的(x,y,w,h,class_id)填充到真实框所属的网格中(对应标记索引)
    网格的置信度设为1
    
    :argument 
        bboxes:预处理后图像ground_truth的[x_min,y_min,x_max,y_max,class_id]
        
    :return
        label_sbbox, label_mbbox:含有标签的ground_truth [x,y,w,h,confidence,class_id_arr]
        其中
        xywh为ground_truth的xywh
        confidence=1.
        class_id_arr 是一个长度为num_class数组，groud_truth的类对应的数组元素为1，数组其余元素为0
        
        sbboxes, mbboxes, lbboxes:不含标签的ground_truth [x,y,w,h]
        其中
        xywh为ground_truth的xywh
        eg:sbboxes[1][1,...]表示在第一种缩放比例下，第1个bbox的[x,y,w,h]
    """
    def preprocess_true_boxes(self, bboxes):
        # 输出图像的尺寸，每个grid的anchor box数，5代表[x, y, w, h, conf, classid]，num_classes存储每个类别的概率
        label = [np.zeros((self.train_output_sizes[i], self.train_output_sizes[i], self.anchor_per_scale,
                           5 + self.num_classes)) for i in range(3)]
        #每种缩放尺度下的ground_truth的xywh
        bboxes_xywh = [np.zeros((self.max_bbox_per_scale, 4)) for _ in range(3)]
        # 不同缩放尺度下的bounding box总数量
        bbox_count = np.zeros((3,))
        # 对每一个选框做处理
        for bbox in bboxes:
            # bbox=[x_min,y_min,x_max,y_max,classid]
            bbox_coor = bbox[:4]
            bbox_class_ind = bbox[4]
            # one_hot，将目标所属类别设为1其余设为0
            onehot = np.zeros(self.num_classes, dtype=np.float)
            onehot[bbox_class_ind] = 1.0
            #one_hot平滑处理的作用？
            uniform_distribution = np.full(self.num_classes, 1.0 / self.num_classes)
            beta = 0.01
            smooth_onehot = onehot * (1 - beta) + beta * uniform_distribution
            #由[x_min,y_min,x_max,y_max]到[x,y,w h]
            bbox_xywh = np.concatenate([(bbox_coor[2:] + bbox_coor[:2]) * 0.5, bbox_coor[2:] - bbox_coor[:2]], axis=-1)
            #np.newaxis的作用就是在这一位置增加一个一维
            #bbox_xywh_scaled为在[8,16,32]三种缩放比例下的x,y,w,h
            bbox_xywh_scaled = 1.0 * bbox_xywh[np.newaxis, :] / self.strides[:, np.newaxis]

            iou = []
            exist_positive = False
            #对每个ground_truth在每种缩放比例下做如下处理
            for i in range(3):
                # 设定变量，用于存储每种缩放比例下3 个 anchor 框的中心位置和宽高
                #anchors_xywh 3*4
                anchors_xywh = np.zeros((self.anchor_per_scale, 4))
                #将这3个anchor box都偏置到网络中心
                anchors_xywh[:, 0:2] = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32) + 0.5
                #填充这3个anchor box的宽高,anchor box的wh是聚类后的9个聚类中心经过reshape后的[i,3,2]
                anchors_xywh[:, 2:4] = self.anchors[i]
                #计算每种缩放比例下ground_truth和三个anchor box的IOU值
                iou_scale = self.bbox_iou(bbox_xywh_scaled[i][np.newaxis, :], anchors_xywh)
                iou.append(iou_scale)
                #找出iou>0.3的选框
                iou_mask = iou_scale > 0.3

                #当ground_truth与三种anchor boxes任意一个IOU>0.3时：
                if np.any(iou_mask):
                    #根据ground_truth的坐标信息计算所属网格左上角的位置
                    #np.floor()向下取整
                    xind, yind = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32)

                    label[i][yind, xind, iou_mask, :] = 0
                    #目标的位置尺寸信息,依旧是未经缩放的ground_truth的xywh
                    label[i][yind, xind, iou_mask, 0:4] = bbox_xywh
                    #有物体
                    label[i][yind, xind, iou_mask, 4:5] = 1.0
                    #类别信息
                    label[i][yind, xind, iou_mask, 5:] = smooth_onehot

                    bbox_ind = int(bbox_count[i] % self.max_bbox_per_scale)
                    bboxes_xywh[i][bbox_ind, :4] = bbox_xywh
                    #记录每种缩放比例下bbox数量
                    bbox_count[i] += 1

                    exist_positive = True

            if not exist_positive:
                #将iou排成一排，选出iou值对大的下标
                best_anchor_ind = np.argmax(np.array(iou).reshape(-1), axis=-1)
                #最佳iou所在的缩放比例
                best_detect = int(best_anchor_ind / self.anchor_per_scale)
                #最佳iou在最佳缩放比例下的第几个anchor box
                best_anchor = int(best_anchor_ind % self.anchor_per_scale)
                #最佳选框的中心位置
                xind, yind = np.floor(bbox_xywh_scaled[best_detect, 0:2]).astype(np.int32)
                #IOU值最大的label
                label[best_detect][yind, xind, best_anchor, :] = 0
                label[best_detect][yind, xind, best_anchor, 0:4] = bbox_xywh
                label[best_detect][yind, xind, best_anchor, 4:5] = 1.0
                label[best_detect][yind, xind, best_anchor, 5:] = smooth_onehot

                bbox_ind = int(bbox_count[best_detect] % self.max_bbox_per_scale)
                bboxes_xywh[best_detect][bbox_ind, :4] = bbox_xywh
                bbox_count[best_detect] += 1

        label_sbbox, label_mbbox, label_lbbox = label

        sbboxes, mbboxes, lbboxes = bboxes_xywh
        #返回在不同尺度下的标签信息[x,y,w,h,1,classindex]，bounding box信息:每个选框(最佳选框或iou>0.3的anchor box)的[x,y,w,h]
        return label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes

    def __len__(self):
        return self.num_batchs
